/*
** Code cache for shared function prototypes.
** Copyright (C) 2005-2025 Mike Pall. See Copyright Notice in luajit.h
*/

#define lj_codecache_c
#define LUA_CORE

#include <stdlib.h>
#include <string.h>
#include "lua.h"
#include "lauxlib.h"
#include "lj_obj.h"
#include "lj_gc.h"
#include "lj_err.h"
#include "lj_debug.h"
#include "lj_func.h"
#include "lj_tab.h"
#include "lj_state.h"
#include "lj_str.h"
#include "lj_codecache.h"

int luaL_loadfilex_nocache(lua_State *L, const char *filename, const char *mode);
static int luajit_clonefunction(lua_State *L, const void *proto);

#define fixobject(o) ((o)->marked |= LJ_GC_WHITE0 | LJ_GC_FIXED | LJ_GC_SFIXED)

/* Fix all strings in a prototype to prevent GC collection */
static void fixproto_strings(GCproto *pt)
{
  MSize i;
  
  /* Fix constant strings in the prototype */
  for (i = 0; i < pt->sizekgc; i++) {
    GCobj *o = proto_kgc(pt, ~(ptrdiff_t)i);
    if (o->gch.gct == ~LJ_TSTR) {
      fixobject(gco2str(o));
    }
  }
  
  /* Fix chunk name string */
  if (pt->chunkname.gcptr64 != 0) {
    GCstr *chunkname = proto_chunkname(pt);
    if (chunkname) {
      fixobject(chunkname);
    }
  }
  
  /* Fix trace if present */
#if LJ_HASJIT
  if (pt->trace) {
    /* Note: trace is a trace number, not a direct GC object pointer.
     * The actual trace object will be marked during normal GC traversal.
     * We don't need to fix it here since it's not directly accessible. */
  }
#endif
  
  /* Note: uvinfo and varinfo contain embedded string data, not GC string objects.
   * These strings are part of the prototype's memory and are automatically 
   * protected by fixing the prototype itself. No additional fixing needed. */
}


void lj_jit_off(lua_State *L) {
  /* 为使用共享 GCproto 的线程禁用 JIT
   * 这确保多线程环境下的完全安全 */
  lua_getglobal(L, "jit");
  if (lua_istable(L, -1)) {
    lua_getfield(L, -1, "off");
    if (lua_isfunction(L, -1)) {
      if (lua_pcall(L, 0, 0, 0) == 0) {
        /* JIT 成功禁用 */
      }
    } else {
      lua_pop(L, 1); /* pop non-function */
    }
  }
  lua_pop(L, 1); /* pop jit table or nil */
}

#if LJ_HASFFI && (LJ_TARGET_POSIX || LJ_TARGET_WINDOWS)
#include <pthread.h>
#define CODECACHE_LOCK_INIT(cc)    do { \
  cc->lock = malloc(sizeof(pthread_mutex_t)); \
  pthread_mutex_init((pthread_mutex_t*)cc->lock, NULL); \
} while(0)
#define CODECACHE_LOCK_FREE(cc)    do { \
  if (cc->lock) { \
    pthread_mutex_destroy((pthread_mutex_t*)cc->lock); \
    free(cc->lock); \
    cc->lock = NULL; \
  } \
} while(0)
#define CODECACHE_LOCK(cc)         pthread_mutex_lock((pthread_mutex_t*)cc->lock)
#define CODECACHE_UNLOCK(cc)       pthread_mutex_unlock((pthread_mutex_t*)cc->lock)
#else
#define CODECACHE_LOCK_INIT(cc)    ((void)0)
#define CODECACHE_LOCK_FREE(cc)    ((void)0)
#define CODECACHE_LOCK(cc)         ((void)0)
#define CODECACHE_UNLOCK(cc)       ((void)0)
#endif

CodeCache lj_codecache_global = {0};

#define CACHE_REGISTRY_KEY "_CODECACHE_"

void lj_codecache_init(void)
{
  CodeCache *cc = &lj_codecache_global;
  if (cc->cache_L) return;
  
  cc->cache_L = luaL_newstate();
  if (cc->cache_L) {
    lua_newtable(cc->cache_L);
    lua_setfield(cc->cache_L, LUA_REGISTRYINDEX, CACHE_REGISTRY_KEY);
    
  }
  
  CODECACHE_LOCK_INIT(cc);
}

// Traverse all GC objects and remove LJ_GC_FIXED flag from prototypes
static void lj_codecache_free_fixed_protos(lua_State *L)
{
  global_State *g = G(L);
  GCRef *p = &g->gc.root;
  GCobj *o;
  
  while ((o = gcref(*p)) != NULL) {
    if (o->gch.gct == ~LJ_TPROTO) {
      GCproto *pt = gco2pt(o);
      if (pt->marked & LJ_GC_FIXED) {
        pt->marked &= ~LJ_GC_FIXED;
      }
    }
    p = &o->gch.nextgc;
  }
  
  lua_gc(L, LUA_GCCOLLECT, 0);
}

void lj_codecache_free(void)
{
  CodeCache *cc = &lj_codecache_global;
  if (!cc->cache_L) return;
  
  CODECACHE_LOCK(cc);
  if (cc->cache_L) {
    lj_codecache_free_fixed_protos(cc->cache_L);
    lua_close(cc->cache_L);
    cc->cache_L = NULL;
  }
  CODECACHE_UNLOCK(cc);
  CODECACHE_LOCK_FREE(cc);
}

int lj_codecache_loadfile(lua_State *L, const char *filename, const char *mode)
{
  CodeCache *cc = &lj_codecache_global;
  if (!cc->cache_L) lj_codecache_init();
  if (!cc->cache_L) return LUA_ERRMEM;
  
  
  lua_State *cache_L = cc->cache_L;
  
  CODECACHE_LOCK(cc);
  
  /* Get cache table */
  lua_getfield(cache_L, LUA_REGISTRYINDEX, CACHE_REGISTRY_KEY);
  if (!lua_istable(cache_L, -1)) {
    lua_pop(cache_L, 1);
    CODECACHE_UNLOCK(cc);
    return LUA_ERRMEM;
  }
  
  lua_pushstring(cache_L, filename);
  lua_rawget(cache_L, -2);  /* cache_table[filename] */
  if (lua_isfunction(cache_L, -1)) {
    GCfunc *fn = funcV(cache_L->top - 1);
    GCproto *pt = NULL;
    if (isluafunc(fn)) {
      pt = funcproto(fn);
    }
    lua_pop(cache_L, 2);
    CODECACHE_UNLOCK(cc);
    
    if (pt) {
      luajit_clonefunction(L, pt);
      return LUA_OK;
    } else {
      lua_pushstring(L, "cached function is not a Lua function");
      return LUA_ERRRUN;
    }
  }
  
  lua_pop(cache_L, 1);  /* Remove nil result */
  
  int status = luaL_loadfilex_nocache(cache_L, filename, mode);
  if (status != LUA_OK) {
    /* Error occurred, copy error message */
    size_t len;
    const char *msg = lua_tolstring(cache_L, -1, &len);
    lua_pop(cache_L, 1);  /* Remove error from cache_L */
    lua_pop(cache_L, 1);  /* Remove cache table */
    CODECACHE_UNLOCK(cc);
    lua_pushlstring(L, msg, len);
    return status;
  }
  
  GCfunc *fn = funcV(cache_L->top - 1);
  GCproto *pt = NULL;
  if (isluafunc(fn)) {
    pt = funcproto(fn);
    fixobject(pt);
    fixproto_strings(pt);
    lua_pushstring(cache_L, filename);
    lua_pushvalue(cache_L, -2);
    lua_rawset(cache_L, -4);
  }
  
  lua_pop(cache_L, 2);
  CODECACHE_UNLOCK(cc);
  
  if (pt) {
    luajit_clonefunction(L, pt);
    return LUA_OK;
  } else {
    lua_pushstring(L, "compiled function is not a Lua function");
    return LUA_ERRRUN;
  }
}

static int luajit_clonefunction(lua_State *L, const void *proto)
{
  if (!proto) {
    return 0;
  }

  GCproto *cache_pt = (GCproto *)proto;
  GCfunc *fn = lj_func_newL_empty(L, cache_pt, tabref(L->env));

  setfuncV(L, L->top, fn);
  incr_top(L);
  lj_jit_off(L);
  lj_gc_check(L);

  return 0;
}

lua_State *get_codecache_state() {
  return lj_codecache_global.cache_L;
}
